import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
import tensorflow as tf
from tensorflow import keras

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"


#文件路径
train_dir = 'G:/TensorFlow/opencv/image/trainfaces/' # 文件夹路径
num_epochs = 20  #设置训练次数
batch_size = 50  #设置最大批处理尺寸


#图片处理
# rescale像素比例缩放， validation_split验证集数据比例
data_gen = keras.preprocessing.image.ImageDataGenerator(rescale=1. / 255, validation_split=0.2)

# target_size=(64, 64) 对象缩放比例
# batch_size 每次读入多少图片
# class_mode 固定字段，会根据路径下的文件做分类
# subset training validation 训练集和验证集
train_generator = data_gen.flow_from_directory(train_dir,
                                               target_size=(64, 64),
                                               batch_size=batch_size,
                                               class_mode='categorical',
                                               subset='training')
validation_generator = data_gen.flow_from_directory(train_dir,
                                                    target_size=(64, 64),
                                                    batch_size=batch_size,
                                                    class_mode='categorical',
                                                    subset='validation')




# train_generator.__getitem__(1)  获得原始数据

labels = train_generator.class_indices # 获得分类对应
labels


#模型构建
# 模型构建
model = keras.Sequential()

# 卷积层1
model.add(keras.layers.Conv2D(64,64,3,input_shape=(64, 64, 3),activation='relu',padding='same'))
# 池化层1
model.add(keras.layers.MaxPool2D(pool_size=(2, 2)))
model.add(keras.layers.Dropout(0.25))

#卷积层2
model.add(keras.layers.Conv2D(128, 128, 3, activation='relu', padding='same'))
# 池化层2
model.add(keras.layers.MaxPool2D(pool_size=(2, 2)))
model.add(keras.layers.Dropout(0.25))

# 平坦层
model.add(keras.layers.Flatten())

# 全连接层1
model.add(keras.layers.Dense(128, activation='relu'))
model.add(keras.layers.Dropout(0.2))
# 全连接层2
model.add(keras.layers.Dense(64, activation='relu'))
model.add(keras.layers.Dropout(0.2))
# 全连接层3
model.add(keras.layers.Dense(32, activation='relu'))
model.add(keras.layers.Dropout(0.2))
# 全连接层4
model.add(keras.layers.Dense(3, activation='sigmoid'))

#模型训练
# 损失函数
model.compile(loss='categorical_crossentropy',
              optimizer='adam',
              metrics=['accuracy'])

# 模型训练
model.fit_generator(
    train_generator,
    steps_per_epoch=5,  # 100
    validation_steps=1,  # 50
    epochs=20,  # 20个周期
    validation_data=validation_generator)

model.save('model_weight.h5')



import matplotlib.pyplot as plt
import numpy as np



test_x, test_y = validation_generator.__getitem__(1)
labels = (train_generator.class_indices)
labels = dict((v, k) for k, v in labels.items())
preds = model.predict(test_x)
print("正确率")
print(preds)

aver=0



